package com.tt.vtg.utils;

import io.jsonwebtoken.Claims;
import io.jsonwebtoken.Jwts;
import io.jsonwebtoken.SignatureAlgorithm;

import java.io.IOException;
import java.io.ObjectInputStream;
import java.security.Key;
import java.security.KeyFactory;
import java.security.NoSuchAlgorithmException;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.interfaces.RSAPrivateKey;
import java.security.interfaces.RSAPublicKey;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.Arrays;
import java.util.Date;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import com.tt.vtg.plain.info.UserInfo;

public class JWTUtils {
	private JWTUtils(){}
	
	private static final String KEY_ALGORITHM = "RSA";
	
	private static final String ISSUER = "tt-vtg";
	
	private static final String PUBLIC_KEY = "publicKey.key";
	private static final String PRIVATE_KEY = "privateKey.key";
	
	private static final Map<String, Key> KEYMAP;
	
	static{
		KEYMAP = loadKeyBytes();
	}
	
	private static Map<String, Key> loadKeyBytes() {
    	ObjectInputStream publicKeyOS = null;
    	ObjectInputStream privateKeyOS = null;
        try {
            publicKeyOS = new ObjectInputStream(ClassLoader.getSystemClassLoader().getResourceAsStream(PUBLIC_KEY));
            privateKeyOS = new ObjectInputStream(ClassLoader.getSystemClassLoader().getResourceAsStream(PRIVATE_KEY));
			
			RSAPublicKey publicKey = (RSAPublicKey) publicKeyOS.readObject();
            RSAPrivateKey privateKey = (RSAPrivateKey) privateKeyOS.readObject();
            
            Map<String, Key> keyMap = new HashMap<String, Key>();
            
            keyMap.put(PUBLIC_KEY, genPublicKey(publicKey.getEncoded()));
            keyMap.put(PRIVATE_KEY, genPrivateKey(privateKey.getEncoded()));
            
            return keyMap;
		} catch (Exception e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		} finally {
			try {
				publicKeyOS.close();
				privateKeyOS.close();
			} catch (IOException e) {
				// TODO Auto-generated catch block
				e.printStackTrace();
			}
		}
        return null;
    }
	
	private static PublicKey genPublicKey(byte[] keyBytes) {
        X509EncodedKeySpec x509EncodedKeySpec = new X509EncodedKeySpec(keyBytes);

        try {
            KeyFactory factory = KeyFactory.getInstance(KEY_ALGORITHM);
            PublicKey publicKey = factory.generatePublic(x509EncodedKeySpec);
            return publicKey;
        } catch (NoSuchAlgorithmException | InvalidKeySpecException e) {
            // TODO Auto-generated catch block
            e.printStackTrace();
        }
        return null;
    }
    
    private static PrivateKey genPrivateKey(byte[] keyBytes) {
        PKCS8EncodedKeySpec pkcs8EncodedKeySpec = new PKCS8EncodedKeySpec(
                keyBytes);
        try {
            KeyFactory factory = KeyFactory.getInstance(KEY_ALGORITHM);
            PrivateKey privateKey = factory
                    .generatePrivate(pkcs8EncodedKeySpec);
            return privateKey;
        } catch (NoSuchAlgorithmException | InvalidKeySpecException e) {
            // TODO Auto-generated catch block
            e.printStackTrace();
        }
        return null;
    }
	
	public static String createJWTByUser(UserInfo userInfo, long ttl){
		return userInfo==null?null:
			Jwts.builder()
			      .setSubject(userInfo.getUserId())
			      .setIssuer(ISSUER)
			      .claim("roles", userInfo.getRoleIds())
			      .claim("perms", userInfo.getPermissionIds())
			      .setExpiration(new Date(System.currentTimeMillis() + (ttl>0?ttl:0)))
			      .signWith(SignatureAlgorithm.RS512, KEYMAP.get(PRIVATE_KEY))
				  .compact();
	}
	
	@SuppressWarnings("unchecked")
	public static UserInfo parseToken(String token){ 
		// throw ExpiredJwtException, if token has been expired.
		Claims claims = Jwts.parser().setSigningKey(KEYMAP.get(PUBLIC_KEY)).parseClaimsJws(token).getBody();
		UserInfo userInfo = new UserInfo();
		userInfo.setUserId(claims.getSubject());
		userInfo.setPermissionIds(claims.get("perms", List.class));
		userInfo.setRoleIds(claims.get("roles", List.class));
		
		return userInfo;
	}
	
	public static void main(String[] args) throws Exception {
		UserInfo u = new UserInfo();
		u.setUserId("test1");
		u.setName("testName");
		u.setRoleIds(Arrays.asList("Role_admin"));
		u.setPermissionIds(Arrays.asList("perm_admin", "perm_del"));
		
		
		String token = JWTUtils.createJWTByUser(u, 60*1000);
//		String token = JWTUtils.createJWTByUser(u, 60*10);
		
		Thread.sleep(1000);
		System.out.println(token);
		
		System.out.println(JWTUtils.parseToken(token).toString());
	}
}
